/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.extensions.gcp.options;



import com.fasterxml.jackson.annotation.JsonIgnore;

import com.google.api.client.http.HttpRequestInitializer;

import com.google.api.client.util.BackOff;

import com.google.api.client.util.Sleeper;

import com.google.api.services.cloudresourcemanager.CloudResourceManager;

import com.google.api.services.cloudresourcemanager.model.Project;

import com.google.api.services.storage.model.Bucket;

import com.google.api.services.storage.model.Bucket.Encryption;

import com.google.auth.Credentials;

import com.google.auth.http.HttpCredentialsAdapter;

import com.google.cloud.hadoop.util.ChainingHttpRequestInitializer;

import com.google.cloud.hadoop.util.ResilientOperation;

import com.google.cloud.hadoop.util.RetryDeterminer;

import com.bff.gaia.unified.sdk.annotations.Experimental;

import com.bff.gaia.unified.sdk.extensions.gcp.auth.CredentialFactory;

import com.bff.gaia.unified.sdk.extensions.gcp.auth.GcpCredentialFactory;

import com.bff.gaia.unified.sdk.extensions.gcp.auth.NullCredentialInitializer;

import com.bff.gaia.unified.sdk.extensions.gcp.storage.PathValidator;

import com.bff.gaia.unified.sdk.extensions.gcp.util.BackOffAdapter;

import com.bff.gaia.unified.sdk.extensions.gcp.util.RetryHttpRequestInitializer;

import com.bff.gaia.unified.sdk.extensions.gcp.util.Transport;

import com.bff.gaia.unified.sdk.extensions.gcp.util.gcsfs.GcsPath;

import com.bff.gaia.unified.sdk.options.*;

import com.bff.gaia.unified.sdk.util.FluentBackoff;

import com.bff.gaia.unified.sdk.util.InstanceBuilder;

import com.bff.gaia.unified.vendor.guava.com.google.common.annotations.VisibleForTesting;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableList;

import com.bff.gaia.unified.vendor.guava.com.google.common.io.Files;

import com.bff.gaia.unified.sdk.options.Default;

import com.bff.gaia.unified.sdk.options.DefaultValueFactory;

import com.bff.gaia.unified.sdk.options.Description;

import com.bff.gaia.unified.sdk.options.ExperimentalOptions;

import com.bff.gaia.unified.sdk.options.PipelineOptions;

import org.joda.time.Duration;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



import javax.annotation.Nullable;

import java.io.File;

import java.io.IOException;

import java.nio.charset.StandardCharsets;

import java.nio.file.FileAlreadyExistsException;

import java.security.GeneralSecurityException;

import java.util.Locale;

import java.util.Map;

import java.util.regex.Matcher;

import java.util.regex.Pattern;



import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions.checkArgument;

import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Strings.isNullOrEmpty;



/**

 * Options used to configure Google Cloud Platform specific options such as the project and

 * credentials.

 *

 * <p>These options defer to the <a

 * href="https://developers.google.com/accounts/docs/application-default-credentials">application

 * default credentials</a> for authentication. See the <a

 * href="https://github.com/google/google-auth-library-java">Google Auth Library</a> for alternative

 * mechanisms for creating credentials.

 */

@Description("Options used to configure Google Cloud Platform project and credentials.")

public interface GcpOptions extends GoogleApiDebugOptions, PipelineOptions {

  /** Project id to use when launching jobs. */

  @Description(

      "Project id. Required when using Google Cloud Platform services. "

          + "See https://cloud.google.com/storage/docs/projects for further details.")

  @Default.InstanceFactory(DefaultProjectFactory.class)

  String getProject();



  void setProject(String value);



  /**

   * GCP <a href="https://developers.google.com/compute/docs/zones" >availability zone</a> for

   * operations.

   *

   * <p>Default is set on a per-service basis.

   */

  @Description(

      "GCP availability zone for running GCP operations. "

          + "Default is up to the individual service.")

  String getZone();



  void setZone(String value);



  /**

   * The class of the credential factory that should be created and used to create credentials. If

   * gcpCredential has not been set explicitly, an instance of this class will be constructed and

   * used as a credential factory.

   */

  @Description(

      "The class of the credential factory that should be created and used to create "

          + "credentials. If gcpCredential has not been set explicitly, an instance of this class will "

          + "be constructed and used as a credential factory.")

  @Default.Class(GcpCredentialFactory.class)

  Class<? extends CredentialFactory> getCredentialFactoryClass();



  void setCredentialFactoryClass(Class<? extends CredentialFactory> credentialFactoryClass);



  /**

   * The credential instance that should be used to authenticate against GCP services. If no

   * credential has been set explicitly, the default is to use the instance factory that constructs

   * a credential based upon the currently set credentialFactoryClass.

   */

  @JsonIgnore

  @Description(

      "The credential instance that should be used to authenticate against GCP services. "

          + "If no credential has been set explicitly, the default is to use the instance factory "

          + "that constructs a credential based upon the currently set credentialFactoryClass.")

  @Default.InstanceFactory(GcpUserCredentialsFactory.class)

  Credentials getGcpCredential();



  void setGcpCredential(Credentials value);



  /** Experiment to turn on the Streaming Engine experiment. */

  String STREAMING_ENGINE_EXPERIMENT = "enable_streaming_engine";



  /** @deprecated Use STREAMING_ENGINE_EXPERIMENT instead. */

  @Deprecated String WINDMILL_SERVICE_EXPERIMENT = "enable_windmill_service";



  @Description(

      "If true will use Streaming Engine.  Defaults to false unless the experiments enable_streaming_engine or enable_windmill_service are set.")

  @Default.InstanceFactory(EnableStreamingEngineFactory.class)

  boolean isEnableStreamingEngine();



  void setEnableStreamingEngine(boolean value);



  /**

   * Attempts to infer the default project based upon the environment this application is executing

   * within. Currently this only supports getting the default project from gcloud.

   */

  class DefaultProjectFactory implements DefaultValueFactory<String> {

    private static final Logger LOG = LoggerFactory.getLogger(DefaultProjectFactory.class);



    @Override

    public String create(PipelineOptions options) {

      try {

        File configFile;

        if (getEnvironment().containsKey("CLOUDSDK_CONFIG")) {

          configFile = new File(getEnvironment().get("CLOUDSDK_CONFIG"), "properties");

        } else if (isWindows() && getEnvironment().containsKey("APPDATA")) {

          configFile = new File(getEnvironment().get("APPDATA"), "gcloud/properties");

        } else {

          // New versions of gcloud use this file

          configFile =

              new File(

                  System.getProperty("user.home"), ".config/gcloud/configurations/config_default");

          if (!configFile.exists()) {

            // Old versions of gcloud use this file

            configFile = new File(System.getProperty("user.home"), ".config/gcloud/properties");

          }

        }

        String section = null;

        Pattern projectPattern = Pattern.compile("^project\\s*=\\s*(.*)$");

        Pattern sectionPattern = Pattern.compile("^\\[(.*)\\]$");

        for (String line : Files.readLines(configFile, StandardCharsets.UTF_8)) {

          line = line.trim();

          if (line.isEmpty() || line.startsWith(";")) {

            continue;

          }

          Matcher matcher = sectionPattern.matcher(line);

          if (matcher.matches()) {

            section = matcher.group(1);

          } else if (section == null || "core".equals(section)) {

            matcher = projectPattern.matcher(line);

            if (matcher.matches()) {

              String project = matcher.group(1).trim();

              LOG.info(

                  "Inferred default GCP project '{}' from gcloud. If this is the incorrect "

                      + "project, please cancel this Pipeline and specify the command-line "

                      + "argument --project.",

                  project);

              return project;

            }

          }

        }

      } catch (IOException expected) {

        LOG.debug("Failed to find default project.", expected);

      }

      // return null if can't determine

      return null;

    }



    /** Returns true if running on the Windows OS. */

    private static boolean isWindows() {

      return System.getProperty("os.name").toLowerCase(Locale.ENGLISH).contains("windows");

    }



    /** Used to mock out getting environment variables. */

    @VisibleForTesting

    Map<String, String> getEnvironment() {

      return System.getenv();

    }

  }



  /**

   * Attempts to load the GCP credentials. See {@link CredentialFactory#getCredential()} for more

   * details.

   */

  class GcpUserCredentialsFactory implements DefaultValueFactory<Credentials> {

    @Override

    public Credentials create(PipelineOptions options) {

      GcpOptions gcpOptions = options.as(GcpOptions.class);

      try {

        CredentialFactory factory =

            InstanceBuilder.ofType(CredentialFactory.class)

                .fromClass(gcpOptions.getCredentialFactoryClass())

                .fromFactoryMethod("fromOptions")

                .withArg(PipelineOptions.class, options)

                .build();

        return factory.getCredential();

      } catch (IOException | GeneralSecurityException e) {

        throw new RuntimeException("Unable to obtain credential", e);

      }

    }

  }



  /** EneableStreamingEngine defaults to false unless one of the two experiments is set. */

  class EnableStreamingEngineFactory implements DefaultValueFactory<Boolean> {

    @Override

    public Boolean create(PipelineOptions options) {

      return ExperimentalOptions.hasExperiment(options, STREAMING_ENGINE_EXPERIMENT)

          || ExperimentalOptions.hasExperiment(options, WINDMILL_SERVICE_EXPERIMENT);

    }

  }



  /**

   * A GCS path for storing temporary files in GCP.

   *

   * <p>Its default to {@link PipelineOptions#getTempLocation}.

   */

  @Description("A GCS path for storing temporary files in GCP.")

  @Default.InstanceFactory(GcpTempLocationFactory.class)

  @Nullable

  String getGcpTempLocation();



  void setGcpTempLocation(String value);



  /** Returns {@link PipelineOptions#getTempLocation} as the default GCP temp location. */

  class GcpTempLocationFactory implements DefaultValueFactory<String> {

    private static final FluentBackoff BACKOFF_FACTORY =

        FluentBackoff.DEFAULT.withMaxRetries(3).withInitialBackoff(Duration.millis(200));

    static final String DEFAULT_REGION = "us-central1";

    static final Logger LOG = LoggerFactory.getLogger(GcpTempLocationFactory.class);



    @Override

    @Nullable

    public String create(PipelineOptions options) {

      String tempLocation = options.getTempLocation();

      if (isNullOrEmpty(tempLocation)) {

        tempLocation =

            tryCreateDefaultBucket(

                options,

                newCloudResourceManagerClient(options.as(CloudResourceManagerOptions.class))

                    .build());

        options.setTempLocation(tempLocation);

      } else {

        try {

          PathValidator validator = options.as(GcsOptions.class).getPathValidator();

          validator.validateOutputFilePrefixSupported(tempLocation);

        } catch (Exception e) {

          throw new IllegalArgumentException(

              String.format(

                  "Error constructing default value for gcpTempLocation: tempLocation is not"

                      + " a valid GCS path, %s. ",

                  tempLocation),

              e);

        }

      }

      return tempLocation;

    }



    /**

     * Creates a default bucket or verifies the existence and proper access control of an existing

     * default bucket. Returns the location if successful.

     */

    @VisibleForTesting

    static String tryCreateDefaultBucket(PipelineOptions options, CloudResourceManager crmClient) {

      GcsOptions gcsOptions = options.as(GcsOptions.class);



      final String projectId = gcsOptions.getProject();

      checkArgument(!isNullOrEmpty(projectId), "--project is a required option.");



      // Look up the project number, to create a default bucket with a stable

      // name with no special characters.

      long projectNumber = 0L;

      try {

        projectNumber = getProjectNumber(projectId, crmClient);

      } catch (IOException e) {

        throw new RuntimeException("Unable to verify project with ID " + projectId, e);

      }

      String region = DEFAULT_REGION;

      if (!isNullOrEmpty(gcsOptions.getZone())) {

        region = getRegionFromZone(gcsOptions.getZone());

      }

      final String bucketName = "dataflow-staging-" + region + "-" + projectNumber;

      LOG.info("No tempLocation specified, attempting to use default bucket: {}", bucketName);

      Bucket bucket =

          new Bucket()

              .setName(bucketName)

              .setLocation(region)

              .setEncryption(new Encryption().setDefaultKmsKeyName(gcsOptions.getDataflowKmsKey()));

      // Always try to create the bucket before checking access, so that we do not

      // race with other pipelines that may be attempting to do the same thing.

      try {

        gcsOptions.getGcsUtil().createBucket(projectId, bucket);

      } catch (FileAlreadyExistsException e) {

        LOG.debug("Bucket '{}'' already exists, verifying access.", bucketName);

      } catch (IOException e) {

        throw new RuntimeException("Unable create default bucket.", e);

      }



      // Once the bucket is expected to exist, verify that it is correctly owned

      // by the project executing the job.

      try {

        long owner = gcsOptions.getGcsUtil().bucketOwner(GcsPath.fromComponents(bucketName, ""));

        checkArgument(

            owner == projectNumber,

            "Bucket owner does not match the project from --project:" + " %s vs. %s",

            owner,

            projectNumber);

      } catch (IOException e) {

        throw new RuntimeException(

            "Unable to determine the owner of the default bucket at gs://" + bucketName, e);

      }

      return "gs://" + bucketName + "/temp/";

    }



    /**

     * Returns the project number or throws an exception if the project does not exist or has other

     * access exceptions.

     */

    private static long getProjectNumber(String projectId, CloudResourceManager crmClient)

        throws IOException {

      return getProjectNumber(

          projectId,

          crmClient,

          BackOffAdapter.toGcpBackOff(BACKOFF_FACTORY.backoff()),

          Sleeper.DEFAULT);

    }



    /**

     * Returns the project number or throws an error if the project does not exist or has other

     * access errors.

     */

    private static long getProjectNumber(

        String projectId, CloudResourceManager crmClient, BackOff backoff, Sleeper sleeper)

        throws IOException {

      CloudResourceManager.Projects.Get getProject = crmClient.projects().get(projectId);

      try {

        Project project =

            ResilientOperation.retry(

                ResilientOperation.getGoogleRequestCallable(getProject),

                backoff,

                RetryDeterminer.SOCKET_ERRORS,

                IOException.class,

                sleeper);

        return project.getProjectNumber();

      } catch (Exception e) {

        throw new IOException("Unable to get project number", e);

      }

    }



    @VisibleForTesting

    static String getRegionFromZone(String zone) {

      String[] zoneParts = zone.split("-", -1);

      checkArgument(zoneParts.length >= 2, "Invalid zone provided: %s", zone);

      return zoneParts[0] + "-" + zoneParts[1];

    }



    /**

     * Returns a CloudResourceManager client builder using the specified {@link

     * CloudResourceManagerOptions}.

     */

    @VisibleForTesting

    static CloudResourceManager.Builder newCloudResourceManagerClient(

        CloudResourceManagerOptions options) {

      Credentials credentials = options.getGcpCredential();

      if (credentials == null) {

        NullCredentialInitializer.throwNullCredentialException();

      }

      return new CloudResourceManager.Builder(

              Transport.getTransport(),

              Transport.getJsonFactory(),

              chainHttpRequestInitializer(

                  credentials,

                  // Do not log 404. It clutters the output and is possibly even required by the

                  // caller.

                  new RetryHttpRequestInitializer(ImmutableList.of(404))))

          .setApplicationName(options.getAppName())

          .setGoogleClientRequestInitializer(options.getGoogleApiTrace());

    }



    private static HttpRequestInitializer chainHttpRequestInitializer(

        Credentials credential, HttpRequestInitializer httpRequestInitializer) {

      if (credential == null) {

        return new ChainingHttpRequestInitializer(

            new NullCredentialInitializer(), httpRequestInitializer);

      } else {

        return new ChainingHttpRequestInitializer(

            new HttpCredentialsAdapter(credential), httpRequestInitializer);

      }

    }

  }



  /**

   * GCP <a href="https://cloud.google.com/kms/">Cloud KMS</a> key for Dataflow pipelines and

   * buckets created by GcpTempLocationFactory.

   */

  @Description(

      "GCP Cloud KMS key for Dataflow pipelines. Also used by gcpTempLocation as the default key "

          + "for new buckets. Key format is: "

          + "projects/<project>/locations/<location>/keyRings/<keyring>/cryptoKeys/<key>")

  @Experimental

  @Nullable

  String getDataflowKmsKey();



  void setDataflowKmsKey(String dataflowKmsKey);

}